#!/usr/bin/env python
"""
Demo: Memory-Augmented Agent with CogneeRetriever and CogneeWriter

Shows the core memory-agent loop where Cognee enriches every conversation
turn with memories from past sessions alongside the live chat history:

1. Pre-seeded memories: facts from previous sessions are loaded into Cognee
   before the conversation starts, simulating persistent long-term memory.
2. Before each turn: CogneeRetriever fetches relevant memories and the
   OutputAdapter prepends them to the live chat history + current user message.
3. The Agent processes the full context (memories + history + user message).
4. After each turn: CogneeWriter stores the agent's output messages so future
   sessions can recall what happened in this one.

Pipeline structure (per turn):

    query ──► CogneeRetriever ──► memories (list[ChatMessage])  ──┐
                                                                   ├──► OutputAdapter ──► agent.messages
    history + user_message (pipeline inputs) ─────────────────────┘
                                                                         │
                                                                      Agent ──► messages, last_message
                                                                                      │
                                                                               CogneeWriter ──► documents_written

NOTE: This demo assumes CogneeRetriever and CogneeWriter have been updated to
work with ChatMessages rather than Documents. See the review comments for the
required changes.

Prerequisites:
    pip install -e "integrations/cognee"

Set your LLM and vector DB keys (Cognee uses them internally):
    export OPENAI_API_KEY="sk-..."
"""

import logging
import os

# Must be set before importing cognee — its setup_logging() reads LOG_LEVEL at
# import time and installs structlog handlers that bypass standard logging config.
os.environ.setdefault("LOG_LEVEL", "WARNING")

import asyncio

import cognee

from haystack import Pipeline

logging.basicConfig(level=logging.WARNING)
from haystack.components.agents import Agent
from haystack.components.converters import OutputAdapter
from haystack.components.generators.chat import OpenAIChatGenerator
from haystack.dataclasses import ChatMessage

from haystack_integrations.components.retrievers.cognee import CogneeRetriever
from haystack_integrations.components.writers.cognee import CogneeWriter
from haystack_integrations.memory_stores.cognee import CogneeMemoryStore

DATASET_NAME = "agent_memory_demo"

# Facts that would have been stored in a previous session.
SEEDED_MEMORIES = [
    "My name is Alice. I'm a senior data scientist at Acme Corp specialising in NLP and knowledge graphs.",
    "My current project is building an internal documentation search system powered by Haystack and Cognee.",
    "My team: Bob is the ML engineer and Carol handles infrastructure.",
    "I prefer concise answers with Python code examples over long prose explanations.",
]

SYSTEM_PROMPT = (
    "You are a helpful assistant with access to a persistent memory of past conversations. "
    "Any system messages at the start of the conversation contain relevant memories retrieved "
    "from previous interactions — use them to personalise your responses and maintain continuity "
    "across turns. "
    "Be concise. Prefer short answers and Python code examples over long prose unless the user "
    "asks for more detail."
)


def build_memory_agent_pipeline(store: CogneeMemoryStore) -> Pipeline:
    """
    Wire the memory-augmented agent pipeline.

    The OutputAdapter merges three inputs into one flat ChatMessage list for
    the Agent:
      - memories:      retrieved from Cognee by CogneeRetriever
      - history:       accumulated chat messages from previous turns
      - user_messages: the current user message

    unsafe=True is required so that the Jinja NativeEnvironment returns the
    concatenated list as a real Python object rather than its string representation.
    """
    pipeline = Pipeline()

    pipeline.add_component("retriever", CogneeRetriever(memory_store=store))
    pipeline.add_component(
        "injector",
        OutputAdapter(
            template="{{ memories + history + user_messages }}",
            output_type=list[ChatMessage],
            unsafe=True,
        ),
    )
    pipeline.add_component(
        "agent",
        Agent(
            chat_generator=OpenAIChatGenerator(model="gpt-4o-mini"),
            system_prompt=SYSTEM_PROMPT,
        ),
    )
    pipeline.add_component("writer", CogneeWriter(dataset_name=store.dataset_name))

    pipeline.connect("retriever.messages", "injector.memories")
    pipeline.connect("injector.output", "agent.messages")
    pipeline.connect("agent.messages", "writer.messages")

    return pipeline


def run_turn(
    pipeline: Pipeline,
    user_text: str,
    history: list[ChatMessage],
) -> str:
    """
    Run one conversation turn, update history in-place, and return the agent reply.

    History is maintained outside the pipeline so it contains only the clean
    user/assistant exchange — not the memory context messages, which are
    injected fresh on every turn.
    """
    result = pipeline.run(
        {
            "retriever": {"query": user_text},
            "injector": {
                "history": history,
                "user_messages": [ChatMessage.from_user(user_text)],
            },
        }
    )
    last_message = result["agent"]["last_message"]
    reply = last_message.text or "(no reply)"

    # Append this turn to history so the next turn has full conversational context.
    history.append(ChatMessage.from_user(user_text))
    history.append(last_message)

    return reply


async def main():
    print("=== Cognee Memory Agent Pipeline Demo ===\n")

    print("Pruning previous data for a clean start...")
    await cognee.prune.prune_data()
    await cognee.prune.prune_system(metadata=True)
    print("Done.\n")

    # -------------------------------------------------------------------------
    # Pre-seed long-term memories from a simulated previous session.
    # In a real application these would already be in Cognee from earlier runs.
    # -------------------------------------------------------------------------
    print("Seeding memories from a previous session...")
    store = CogneeMemoryStore(search_type="GRAPH_COMPLETION", dataset_name=DATASET_NAME)
    store.add_memories(messages=[ChatMessage.from_user(fact) for fact in SEEDED_MEMORIES])
    print(f"  {len(SEEDED_MEMORIES)} facts stored.\n")

    pipeline = build_memory_agent_pipeline(store)
    history: list[ChatMessage] = []

    # -------------------------------------------------------------------------
    # Multi-turn conversation. The Agent receives retrieved memories + live
    # history + the current message on every turn, so it can both recall past
    # facts and follow the thread of the current exchange.
    # -------------------------------------------------------------------------
    turns = [
        # Turn 1 — recall a specific memory (no history yet)
        "Hi! Can you remind me what project I'm currently working on?",
        # Turn 2 — follow-up that requires turn 1 to make sense
        "What's the tech stack we're using for it?",
        # Turn 3 — recall a different memory fact
        "Who else is on my team, and what are their roles?",
        # Turn 4 — requires both a recalled preference *and* conversation context
        "Based on what you know about me, give me a quick tip for structuring a new Haystack pipeline component.",
    ]

    for user_text in turns:
        print(f"User:  {user_text}")
        reply = run_turn(pipeline, user_text, history)
        print(f"Agent: {reply}\n")

    print("=== Done ===")


if __name__ == "__main__":
    asyncio.run(main())